import chromadb
from llama_index.core import StorageContext, VectorStoreIndex
from llama_index.vector_stores.chroma import ChromaVectorStore


class ChromaTool:

    def __init__(self,chroma_db_path,correction_name):
        self.chroma_db_path = chroma_db_path
        self.correction_name = correction_name

    def getVectorStoreIndex(self):
        # 初始化Chroma客户端
        chroma_db = chromadb.PersistentClient(self.chroma_db_path)

        # 获取 collection
        chroma_correction = chroma_db.get_or_create_collection(self.correction_name)

        # 将chroma 分配到上下文中的VectorStore
        vector_store = ChromaVectorStore(chroma_collection=chroma_correction)
        storage_context = StorageContext.from_defaults(vector_store=vector_store)

        # 加载索引
        return VectorStoreIndex.from_vector_store(vector_store, storage_context=storage_context)
